from torch.distributions.categorical import Categorical
import torch

# 导入工具函数
from Utils.tensor_ops import repeat_at, _2tensor
from torch.distributions import kl_divergence

# 常量定义
EPS = 1e-9  # 防止除零错误的小量
yita = p_hit = 0.14  # 命中概率和探索率参数


def random_process(probs, rsn_flag):
    """
    随机处理函数：实现ε-贪婪策略的变体

    这个函数实现了强化学习中常用的探索策略，通过控制最大概率动作被选择的概率，
    来平衡探索和利用。与标准的ε-贪婪不同，这个函数会降低其他动作的概率
    来补偿不选择最大概率动作的影响。

    Args:
        probs (torch.Tensor): 动作概率分布，形状为(*, n_actions)
        rsn_flag (bool or torch.Tensor): 是否使用随机策略标志
                                       True表示选择最优动作，False表示进行探索

    Returns:
        torch.Tensor: 选择的动作索引，形状为(*,)

    算法说明：
    1. 如果rsn_flag为True，直接选择概率最大的动作（利用）
    2. 如果rsn_flag为False，执行探索：
       - 将最大概率动作的概率调整为p_hat
       - 其他动作的概率按比例缩放
       - 从调整后的分布中采样
    """
    with torch.no_grad():  # 不计算梯度，因为是决策过程
        # 找到概率最大的动作
        max_place = probs.argmax(-1, keepdims=True)
        # 创建最大概率位置的掩码
        mask_max = torch.zeros_like(probs).scatter_(-1, max_place, 1).bool()
        # 获取最大概率值
        pmax = probs[mask_max]

        if rsn_flag:
            # 利用阶段：直接选择最大概率的动作
            assert max_place.shape[-1] == 1
            return max_place.squeeze(-1)
        else:
            # 探索阶段：调整概率分布
            # 计算调整后的最大概率 p_hat
            # 公式：p_hat = pmax + (pmax-1)/(1/yita-1)
            # 这个公式确保其他动作的概率被适当缩放
            p_hat = pmax + (pmax-1)/(1/yita-1)
            # 计算缩放因子 k
            k = 1/(1-yita)
            # 缩放所有概率
            probs *= k
            # 将最大概率设置为计算后的值
            probs[mask_max] = p_hat

            # 从调整后的分布中采样
            dist = Categorical(probs=probs)
            samp = dist.sample()
            assert samp.shape[-1] != 1
            return samp


def random_process_allow_big_yita(probs, rsn_flag):
    """
    允许大yita值的随机处理函数

    这是random_process的改进版本，支持更大的yita值。
    当yita较大时，会进行概率裁剪以避免数值不稳定。

    Args:
        probs (torch.Tensor): 动作概率分布
        rsn_flag (bool or torch.Tensor): 随机策略标志

    Returns:
        torch.Tensor: 选择的动作索引
    """
    with torch.no_grad():
        max_place = probs.argmax(-1, keepdims=True)
        mask_max = torch.zeros_like(probs).scatter_(-1, max_place, 1).bool()
        pmax = probs[mask_max].reshape(max_place.shape)

        if rsn_flag:
            assert max_place.shape[-1] == 1
            return max_place.squeeze(-1)
        else:
            # 创建yita数组，支持向量化的yita值
            yita_arr = torch.ones_like(pmax) * yita
            # 将yita裁剪到不超过最大概率
            yita_arr_clip = torch.minimum(pmax, yita_arr)
            # 计算调整后的最大概率
            p_hat = (pmax - yita_arr_clip) / ((1 - yita_arr_clip) + EPS)
            # 计算缩放因子
            k = 1 / ((1 - yita_arr_clip) + EPS)
            # 缩放概率
            probs *= k
            # 设置最大概率
            probs[mask_max] = p_hat.reshape(-1)

            # 采样
            dist = Categorical(probs=probs)
            samp = dist.sample()
            assert samp.shape[-1] != 1
            return samp


def random_process_with_clamp3(probs, yita, yita_min_prob, rsn_flag):
    """
    带概率裁剪的随机处理函数

    这个函数在探索时会确保最小概率不低于yita_min_prob，
    提供更稳定的探索策略。

    Args:
        probs (torch.Tensor): 动作概率分布
        yita (float): 探索率参数
        yita_min_prob (float): 最小概率阈值
        rsn_flag (bool or torch.Tensor): 随机策略标志

    Returns:
        torch.Tensor: 选择的动作索引
    """
    with torch.no_grad():
        max_place = probs.argmax(-1, keepdims=True)
        mask_max = torch.zeros_like(probs).scatter_(dim=-1, index=max_place, value=1).bool()
        pmax = probs[mask_max].reshape(max_place.shape)

        # 最大概率对应的动作
        act_max = max_place.squeeze(-1)

        # 计算调整后的最大概率
        yita_arr = torch.ones_like(pmax) * yita
        p_hat = (pmax - yita_arr) / ((1 - yita_arr) + EPS)
        # 裁剪到最小概率
        p_hat = p_hat.clamp(min=yita_min_prob)

        # 计算缩放因子，确保概率和为1
        k = (1 - p_hat) / ((1 - pmax) + EPS)
        probs *= k
        probs[mask_max] = p_hat.reshape(-1)

        # 采样
        dist = Categorical(probs=probs)
        act_samp = dist.sample()

        # 根据rsn_flag选择动作
        hit_e = _2tensor(rsn_flag)
        return torch.where(hit_e, act_max, act_samp)


class CCategorical():
    """
    约束分类分布类

    这个类封装了Categorical分布，添加了特殊的采样策略。
    主要用于强化学习中的动作选择，支持自定义的探索策略。

    特点：
    - 继承了PyTorch的Categorical分布的功能
    - 添加了约束采样策略
    - 支持动态调整探索参数
    """

    def __init__(self, planner):
        """
        初始化约束分类分布

        Args:
            planner: 规划器对象，必须包含yita和yita_min_prob属性
                     这两个参数控制探索策略的行为
        """
        self.planner = planner
        # 初始化时不执行任何操作

    def sample(self, dist, eprsn):
        """
        从分布中采样动作

        Args:
            dist (Categorical): PyTorch的分类分布对象
            eprsn (bool or torch.Tensor): 探索标志
                                        True表示利用，False表示探索

        Returns:
            torch.Tensor: 采样的动作索引
        """
        # 获取概率分布的副本
        probs = dist.probs.clone()
        # 使用带裁剪的随机处理函数
        return random_process_with_clamp3(probs, self.planner.yita, self.planner.yita_min_prob, eprsn)

    def register_rsn(self, rsn_flag):
        """
        注册随机标志

        Args:
            rsn_flag: 随机标志，控制是否进行探索
        """
        self.rsn_flag = rsn_flag

    def feed_logits(self, logits):
        """
        从logits创建分类分布

        Args:
            logits (torch.Tensor): 未归一化的对数概率

        Returns:
            Categorical: 分类分布对象

        注意：
        - 如果logits包含无效值（如NaN或无穷大），会打印错误信息
        - 这个方法主要用于从神经网络输出创建分布
        """
        try:
            return Categorical(logits=logits)
        except:
            print('error：无法从logits创建分类分布')


# 使用示例和说明：
"""
CCategorical类的使用示例：

1. 创建规划器：
class MyPlanner:
    def __init__(self):
        self.yita = 0.1  # 探索率
        self.yita_min_prob = 0.01  # 最小概率

planner = MyPlanner()

2. 创建约束分类分布：
ccat = CCategorical(planner)

3. 从神经网络输出采样：
# 假设这是神经网络的输出
logits = torch.tensor([[1.0, 2.0, 3.0], [0.5, 1.5, 2.5]])

# 创建分布
dist = ccat.feed_logits(logits)

# 采样动作
eprsn = False  # 进行探索
actions = ccat.sample(dist, eprsn)
print(actions)  # 可能输出 tensor([2, 2])

4. 控制探索vs利用：
# 利用阶段：选择最优动作
actions_best = ccat.sample(dist, eprsn=True)
print(actions_best)  # 可能输出 tensor([2, 2])

# 探索阶段：随机探索
actions_explore = ccat.sample(dist, eprsn=False)
print(actions_explore)  # 可能输出 tensor([1, 0])

算法原理说明：

1. ε-贪婪策略的变体：
   - 标准ε-贪婪：以ε概率随机选择，以1-ε概率选择最优
   - 本实现：调整概率分布，确保探索时仍然考虑动作的原始概率

2. 概率调整公式：
   - p_hat = (pmax - yita) / (1 - yita)
   - k = (1 - p_hat) / (1 - pmax)
   - 新概率 = 原概率 * k，最大概率除外

3. 数值稳定性：
   - 使用EPS防止除零
   - 支持概率裁剪避免数值不稳定
   - 处理边界情况（如yita > pmax）

应用场景：
- 强化学习中的动作选择
- 需要控制探索-利用平衡的策略
- 概率分布的约束采样
"""